import os
import math
import shutil
import hashlib
from typing import List, Tuple

import imagehash
import pillow_heif
import pyexiv2

from loguru import logger
from imagehash import ImageHash
from PIL import Image
from tqdm import tqdm

# 注册 heif 格式支持
pillow_heif.register_heif_opener()

# 将最大像素限制调高，比如设为 200,000,000
Image.MAX_IMAGE_PIXELS = None


class TidyImage(object):
    def __init__(self, file_path: str, image_hash: ImageHash, md5hash, width: int, height: int, exif=None):
        self.file_path = file_path
        self.phash = image_hash
        self.md5hash = md5hash
        self.width = width
        self.height = height
        self.exif = exif

    @property
    def ratio(self):
        """图片比例"""
        # 获取宽高的最大公约数，简化宽高比
        gcd = math.gcd(self.width, self.height)
        return f"{self.width // gcd}:{self.height // gcd}"

    @property
    def shape(self):
        """图片形状"""
        ratio = self.width / self.height

        # 根据比例和具体的宽高值判断图片形状
        if 1.2 < ratio < 2.5:
            return '横图'
        elif 0.8 > ratio > 0.4:
            return '竖图'
        elif 0.8 <= ratio <= 1.2:
            return '方形'
        elif ratio >= 2.5:
            return '细长横图'
        elif ratio <= 0.4:
            return '细长竖图'
        else:
            return '奇形怪状'


def hamming_distance(hash1, hash2):
    """汉明距离: 比较两个 pHash 值的二进制表示，通过 XOR 运算符计算两个哈希值之间的汉明距离，并统计不同的位数。"""
    return bin(int(hash1) ^ int(hash2)).count('1')


def scan_lib(lib_dir, use_image_hash=False, image_hashing='phash', use_md5=False):
    """
    扫描图片库

    :param lib_dir:
    :param use_image_hash: 启用计算图片hash
    :param image_hashing: 选择哈希算法，可选项 ['phash', 'whash']
    :param use_md5: 计算文件MD5
    :return:
    """
    logger.info(f'开始扫描图片库: {lib_dir}')

    images = []

    total_files = len([file for file in os.listdir(lib_dir)])

    with tqdm(
            total=total_files,
            ncols=100,
            desc="🔍 Scanning images",
            # 可选颜色 [hex (#00ff00), BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE]
            colour='green',
            unit="file",
            # 格式化: 使用 bar_format 来自定义进度条的格式，可以显示已处理文件数量、总文件数量、处理速度等。
            # {l_bar} 是进度条的左侧部分，包含描述和百分比。
            # {bar} 是实际的进度条。
            # {n_fmt}/{total_fmt} 显示当前进度和总数。
            # {rate_fmt} 显示处理速率。
            bar_format='{l_bar}{bar} | {n_fmt}/{total_fmt} files [{rate_fmt}]',
    ) as pbar:
        for file in os.listdir(lib_dir):
            image_path = os.path.join(lib_dir, file)
            # print(file_path)

            with Image.open(image_path) as img:
                md5 = None
                if use_md5:
                    hasher = hashlib.md5()
                    hasher.update(img.tobytes())
                    md5 = hasher.hexdigest()

                img_hash = None
                if use_image_hash:
                    if image_hashing == 'phash':
                        img_hash = imagehash.phash(img)
                    elif image_hashing == 'whash':
                        img_hash = imagehash.whash(img)

                try:
                    exif_data = img.getexif()
                except Exception as e:
                    exif_data = None
                    logger.error(f'读取 {file} EXIF信息失败: {e}')
                    logger.debug(f'清除 {file} EXIF信息')
                    pyexiv2.ImageData(img.tobytes()).clear_exif()

                tidy_img = TidyImage(
                    image_path, image_hash=img_hash, md5hash=md5, width=img.width, height=img.height, exif=exif_data
                )
                images.append(tidy_img)

            # Update progress bar for each file processed
            pbar.update(1)

    logger.info(f'扫描图片库完成，图片数量：{len(images)}')
    return images


def find_similar_images(images: List[TidyImage], threshold=10) -> List[Tuple[TidyImage, TidyImage, float]]:
    similar_images = []

    # 计算总的比较次数
    total_comparisons = (len(images) * (len(images) - 1)) // 2

    with tqdm(
            total=total_comparisons,
            ncols=100,
            desc="finding similar images",
            unit="file",
            colour='yellow'
    ) as pbar:
        for i in range(len(images)):
            for j in range(i + 1, len(images)):  # 只与后面的图片对比，避免重复
                image1 = images[i]
                image2 = images[j]

                # 计算汉明距离
                distance = math.fabs(image1.phash - image2.phash)

                # 判断距离是否小于阈值，若小于则认为相似
                if distance <= threshold:
                    similar_images.append((image1, image2, distance))

                pbar.update(1)

    return similar_images


def organize(images: List[TidyImage], output_dir, prefix):
    """整理图片"""

    similar_dir = os.path.join(output_dir, f'{prefix}-相似图片')
    os.makedirs(similar_dir, exist_ok=True)
    similar_images = find_similar_images(images)

    with tqdm(total=len(similar_images), ncols=100, desc="Copying similar images", unit="file", colour='CYAN') as pbar:
        for index, (img1, img2, distance) in enumerate(similar_images):
            group_dir = os.path.join(similar_dir, f'重复-{index}-差异度-{distance}')
            os.makedirs(group_dir, exist_ok=True)
            shutil.copy(img1.file_path, os.path.join(group_dir, os.path.basename(img1.file_path)))
            shutil.copy(img2.file_path, os.path.join(group_dir, os.path.basename(img2.file_path)))

            pbar.update(1)

    with tqdm(total=len(images), ncols=100, desc="Copying images", unit="file", colour='blue') as pbar:
        for image in images:
            target_dir = os.path.join(output_dir, f'{prefix}-{image.shape}')
            os.makedirs(target_dir, exist_ok=True)
            shutil.copy(image.file_path, os.path.join(target_dir, os.path.basename(image.file_path)))

            pbar.update(1)


if __name__ == '__main__':
    data = scan_lib(r'C:\Users\deali\Pictures\手机图片\壁纸 星穹铁道', use_image_hash=True, image_hashing='phash')
    organize(data, r'C:\Users\deali\Pictures\image-lib', '壁纸-星穹铁道')
